from tensorboard.backend.event_processing import event_accumulator
import pandas as pd

# from tensorflow.python.summary.summary_iterator import summary_iterator

# 指定日志文件路径
# log_path = "output/da_aug"
log_path = "output/da_aug/bsz8/conll04/remove_triple/t5-base/RE_STRICT/0906_16_15/runs/Sep06_16-15-14_ccip33"

import os

save_dir = "./export/eval"
os.makedirs(save_dir, exist_ok=True)


# 创建事件累加器
def export_last_step_eval_result(log_file):
    ea = event_accumulator.EventAccumulator(log_file)
    ea.Reload()

    # 获取所有标量数据的标签
    tags = ea.Tags()["scalars"]

    # 要导出的tag: tag 名称中包括 precision, recall, f1, loss
    export_tags = list(
        filter(
            lambda tag: "precision" in tag or "recall" in tag or "f1-score" in tag, tags
        )
    )

    result_df = pd.DataFrame()
    for tag in export_tags:
        events = ea.Scalars(tag)
        df = pd.DataFrame(events)

        df.columns = ["wall_time", "step", tag]

        # 只保留最后一个step的数据
        last_step = df.iloc[-1]
        # 创建一个新的DataFrame，只包含最后一步的结果
        last_step_df = pd.DataFrame(
            {
                "tag": [tag],
                # value值保留4位小数，四舍五入
                "value": [round(last_step[tag], 4)],
            }
        )

        # 将这个标签的最后一步结果添加到总结果中
        result_df = pd.concat([result_df, last_step_df], ignore_index=True)

    # 保存所有标签的最后一步结果到一个CSV文件

    if result_df.empty:
        return

    file_name = "_".join(log_file.split("/")[2:])

    output_file = f"{save_dir}/{file_name}.csv"
    result_df.to_csv(output_file, index=False)
    print(f"Exported last step of all tags to {output_file}")


if __name__ == "__main__":
    # 获取所有日志文件
    # log_files = [
    #     os.path.join(log_path, file)
    #     for file in os.listdir(log_path)
    #     if file.endswith("tfevents")
    # ]

    # for log_file in log_files:
    #     export_last_step_eval_result(log_file)

    # export_last_step_eval_result(log_path)
    output_dir = "output/da_aug"

    def find_ccip33_in_runs(output_dir):
        ccip33_folders = []

        # 使用 os.walk 遍历目录
        for root, dirs, files in os.walk(output_dir):
            print(root, dirs, files)
            # 检查当前目录名是否为 "runs"
            if os.path.basename(root) == "runs":
                # 在 "runs" 目录中查找包含 "ccip_33" 的子目录
                for dir in dirs:
                    if "ccip33" in dir:
                        full_path = os.path.join(root, dir)
                        ccip33_folders.append(full_path)

        return ccip33_folders

    # 找到该文件夹下所有名为runs 的文件夹

    # print(find_runs_folders(output_dir))

    log_files = find_ccip33_in_runs(output_dir)

    print(log_files)

    for log_file in log_files:
        print(log_file)
        for file in os.listdir(log_file):
            if "tfevents" in file:
                export_last_step_eval_result(os.path.join(log_file, file))
